#pragma once
#include <vector>
#include <fmt/color.h>
#include <fmt/format.h>
#include <glm/glm.hpp>
#include <glm/gtc/type_ptr.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include "../opengl/opengl.hpp"
#include "./stb_image.h"
#include <assimp/Importer.hpp>
#include <assimp/scene.h>
#include <assimp/postprocess.h>

namespace Model {

using namespace glm;
namespace gl = opengl;

struct Vertex {
    vec3 position;
    vec3 normal;
    vec2 texCoord;
};

struct Mesh {
public:
    std::vector<Vertex> vertexData{};
    std::vector<gl::u32> indexData{};
    std::vector<gl::texture2d> diffuseMaps{}, specularMaps{}, normalMaps{}, heightMaps{};

public:
    gl::buffer *vbo = nullptr;
    gl::buffer *ibo = nullptr;
    gl::buffer numTexturesUBO;

public:
    Mesh() : numTexturesUBO(16) {}
    Mesh(const Mesh &) = delete;
    Mesh &operator=(const Mesh &) = delete;
    Mesh(Mesh &&other)
        : vertexData(std::move(other.vertexData)), indexData(std::move(other.indexData)),
          diffuseMaps(std::move(other.diffuseMaps)), specularMaps(std::move(other.specularMaps)),
          normalMaps(std::move(other.normalMaps)), heightMaps(std::move(other.heightMaps)), vbo(other.vbo),
          ibo(other.ibo), numTexturesUBO(std::move(other.numTexturesUBO)) {
        other.vbo = nullptr;
        other.ibo = nullptr;
    }
    Mesh &operator=(Mesh &&other) {
        vertexData = std::move(other.vertexData);
        indexData = std::move(other.indexData);
        diffuseMaps = std::move(other.diffuseMaps);
        specularMaps = std::move(other.specularMaps);
        normalMaps = std::move(other.normalMaps);
        heightMaps = std::move(other.heightMaps);
        vbo = other.vbo;
        ibo = other.ibo;
        other.vbo = nullptr;
        other.ibo = nullptr;
        numTexturesUBO = std::move(other.numTexturesUBO);
        return *this;
    }

    void create() {
        vbo = new gl::buffer(vertexData.size() * sizeof(Vertex), vertexData.data());
        ibo = new gl::buffer(indexData.size() * sizeof(indexData[0]), indexData.data());
    }
    void render() const {
        gl::u32 s1 = diffuseMaps.size(), s2 = specularMaps.size(), s3 = normalMaps.size(), s4 = heightMaps.size();
        gl::u32 TextureBlock[4] = {s1, s2, s3, s4};
        numTexturesUBO.set_data(TextureBlock, 16);
        numTexturesUBO.bind_to(gl::buffer::uniform_buffer, 1);
        for (gl::u32 i = 0; i < s1; ++i) {
            diffuseMaps[i].bind_to(i);
        }
        for (gl::u32 i = 0; i < s2; ++i) {
            specularMaps[i].bind_to(i);
        }
        for (gl::u32 i = 0; i < s3; ++i) {
            normalMaps[i].bind_to(i);
        }
        for (gl::u32 i = 0; i < s4; ++i) {
            heightMaps[i].bind_to(i);
        }
        gl::api::glDrawElements(gl::mode_triangles, indexData.size(), gl::type_u32, nullptr);
    }
    ~Mesh() {
        if (vbo) {
            delete vbo;
        }
        if (ibo) {
            delete ibo;
        }
    }
};

struct Model {
public:
    Model() {}
    Model(const std::string &filePath) { loadFile(filePath); }

public:
    void loadFile(const std::string &filePath) {
        using namespace Assimp;
        Importer importer;
        constexpr auto flags =
            aiProcess_Triangulate | aiProcess_FlipUVs | aiProcess_GenNormals | aiProcess_OptimizeMeshes;
        auto scene = importer.ReadFile(filePath, flags);
        if (!scene || scene->mFlags & AI_SCENE_FLAGS_INCOMPLETE || !scene->mRootNode) {
            fmt::print(fmt::fg(fmt::color::red), "[Assimp error] {}\n", importer.GetErrorString());
            return;
        }
        fileDir = filePath.substr(0, filePath.find_last_of('/'));
        processNode(scene->mRootNode, scene);
    }
    void render() const {
        vao.bind();
        program.use();
        for (const auto &mesh : meshes) {
            vao.set_vertex_buffer(*mesh.vbo, sizeof(Vertex));
            vao.set_index_buffer(*mesh.ibo);
            mesh.render();
        }
    }

private:
    // void processMesh(const aiMesh* mesh){}
    void addMesh(const aiMesh *mesh, const aiScene *scene) {
        meshes.push_back(Mesh{});
        Mesh *m = &meshes.back();
        if (auto texCoords = mesh->mTextureCoords[0]; texCoords != nullptr) {
            for (gl::u32 i = 0; i < mesh->mNumVertices; ++i) {
                auto v = mesh->mVertices[i];
                auto n = mesh->mNormals[i];
                auto c = texCoords[i];
                m->vertexData.push_back(Vertex{{v.x, v.y, v.z}, {n.x, n.y, n.z}, {c.x, c.y}});
            }
        } else {
            for (gl::u32 i = 0; i < mesh->mNumVertices; ++i) {
                auto v = mesh->mVertices[i];
                auto n = mesh->mNormals[i];
                m->vertexData.push_back(Vertex{{v.x, v.y, v.z}, {n.x, n.y, n.z}, {0.0f, 0.0f}});
            }
        }
        for (gl::u32 i = 0; i < mesh->mNumFaces; ++i) {
            auto face = mesh->mFaces[i];
            for (gl::u32 j = 0; j < face.mNumIndices; ++j) {
                m->indexData.push_back(face.mIndices[j]);
            }
        }
        // Materials
        auto mat = scene->mMaterials[mesh->mMaterialIndex];
        // Diffuse texture
        auto textureFromFile = [&](const char *fileName) {
            auto path = fmt::format("{}/{}", fileDir, fileName);
            int w, h, n;
            auto data = stbi_load(path.c_str(), &w, &h, &n, 4);
            gl::texture2d texture;
            if (data != nullptr) {
                texture.create(gl::image2d{w, h, gl::format_rgba8, data});
            }
            stbi_image_free(data);
            return std::move(texture);
        };
        aiString str;
        for (gl::u32 i = 0; i < mat->GetTextureCount(aiTextureType_DIFFUSE); ++i) {
            mat->GetTexture(aiTextureType_DIFFUSE, i, &str);
            m->diffuseMaps.push_back(textureFromFile(str.C_Str()));
        }
        for (gl::u32 i = 0; i < mat->GetTextureCount(aiTextureType_SPECULAR); ++i) {
            mat->GetTexture(aiTextureType_SPECULAR, i, &str);
            m->specularMaps.push_back(textureFromFile(str.C_Str()));
        }
        for (gl::u32 i = 0; i < mat->GetTextureCount(aiTextureType_HEIGHT); ++i) {
            mat->GetTexture(aiTextureType_HEIGHT, i, &str);
            m->normalMaps.push_back(textureFromFile(str.C_Str()));
        }
        for (gl::u32 i = 0; i < mat->GetTextureCount(aiTextureType_AMBIENT); ++i) {
            mat->GetTexture(aiTextureType_DIFFUSE, i, &str);
            m->heightMaps.push_back(textureFromFile(str.C_Str()));
        }
        m->create();
    }
    void processNode(const aiNode *node, const aiScene *scene) {
        for (gl::u32 i = 0; i < node->mNumMeshes; ++i) {
            auto mesh = scene->mMeshes[node->mMeshes[i]];
            addMesh(mesh, scene);
        }
        for (gl::u32 i = 0; i < node->mNumChildren; i++) {
            processNode(node->mChildren[i], scene);
        }
    }

public:
    std::string fileDir{};
    std::vector<Mesh> meshes{};
    gl::vertex_array vao{
        gl::vertex_attribute(0, gl::type_f32, 3, 0),  // position
        gl::vertex_attribute(1, gl::type_f32, 3, 12), // normal
        gl::vertex_attribute(2, gl::type_f32, 2, 24), // texCoord
    };
    gl::program program{
        gl::shader::from_file(gl::shader::vertex_shader, "../shader/model/vs.glsl"),
        gl::shader::from_file(gl::shader::fragment_shader, "../shader/model/fs.glsl"),
    };
};

}; // namespace Model